package com.asiainfo.sqlextractor.exec;

import com.alibaba.fastjson.JSONObject;
import com.asiainfo.sqlextractor.adapter.MemSchema;
import com.asiainfo.sqlextractor.adapter.MemTable;
import com.asiainfo.sqlextractor.adapter.TableDefine;
import com.asiainfo.sqlextractor.constant.Def;
import com.asiainfo.sqlextractor.functions.Functions;
import com.asiainfo.sqlextractor.functions.MemSqlOperatorTable;
import com.asiainfo.sqlextractor.functions.StandardConvertletTableExt;
import com.asiainfo.sqlextractor.util.JavaTypeResolve;
import com.google.common.collect.ImmutableList;
import org.apache.calcite.DataContext;
import org.apache.calcite.adapter.enumerable.EnumerableConvention;
import org.apache.calcite.adapter.enumerable.EnumerableInterpretable;
import org.apache.calcite.adapter.enumerable.EnumerableRel;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.adapter.java.ReflectiveSchema;
import org.apache.calcite.config.CalciteConnectionConfig;
import org.apache.calcite.config.CalciteConnectionConfigImpl;
import org.apache.calcite.config.Lex;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.QueryProvider;
import org.apache.calcite.linq4j.tree.Types;
import org.apache.calcite.plan.*;
import org.apache.calcite.plan.RelOptTable.ViewExpander;
import org.apache.calcite.plan.volcano.VolcanoPlanner;
import org.apache.calcite.prepare.CalciteCatalogReader;
import org.apache.calcite.rel.*;
import org.apache.calcite.rel.core.Sort;
import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider;
import org.apache.calcite.rel.metadata.RelMetadataQueryBase;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexExecutorImpl;
import org.apache.calcite.runtime.Bindable;
import org.apache.calcite.schema.ScalarFunction;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.schema.impl.ScalarFunctionImpl;
import org.apache.calcite.sql.*;
import org.apache.calcite.sql.fun.SqlLibrary;
import org.apache.calcite.sql.fun.SqlLibraryOperatorTableFactory;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.parser.SqlParser.Config;
import org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.calcite.sql.parser.impl.SqlParserImpl;
import org.apache.calcite.sql.type.*;
import org.apache.calcite.sql.util.SqlOperatorTables;
import org.apache.calcite.sql.validate.SqlConformanceEnum;
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
import org.apache.calcite.sql.validate.SqlValidator;
import org.apache.calcite.sql.validate.SqlValidatorUtil;
import org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.calcite.tools.Programs;
import org.apache.calcite.util.ImmutableIntList;
import org.apache.calcite.util.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;

import static org.apache.calcite.sql.parser.SqlParser.DEFAULT_IDENTIFIER_MAX_LENGTH;

public final class SqlExtractor {

    private static final Logger LOGGER = LoggerFactory.getLogger(SqlExtractor.class);

    private static final String DEFAULT_SCHEMA = "JSON";

    private final RelDataTypeFactory typeFactory;

    private final CalciteSchema rootSchema;

    private final CalciteCatalogReader catalogReader;

    private final Config parserConfig;

    private final SqlValidator validator;

    private final SqlToRelConverter relConverter;

    private final String sql;

    private RelNode bestPlan;

    private RelDataType rowType;

    private Bindable<Object[]> executablePlan;

    public SqlExtractor(final String sql, final Properties props) throws SqlParseException {
        this.sql = sql;
        CalciteConnectionConfig config = new CalciteConnectionConfigImpl(props);
        this.typeFactory = new JavaTypeFactoryImpl();
        this.rootSchema = CalciteSchema.createRootSchema(false, false);

        registerTables(this.rootSchema);

        registerFunctions(this.rootSchema);

        this.catalogReader = new CalciteCatalogReader(rootSchema, Collections.singletonList(DEFAULT_SCHEMA),
                typeFactory, config);

        this.parserConfig = SqlParser.config()
                .withLex(Lex.ORACLE)
                .withIdentifierMaxLength(DEFAULT_IDENTIFIER_MAX_LENGTH)
                .withConformance(SqlConformanceEnum.ORACLE_12)
                .withParserFactory(SqlParserImpl.FACTORY);

        MemSqlOperatorTable operatorTable = new MemSqlOperatorTable(
                SqlOperatorTables.chain(SqlStdOperatorTable.instance(),
                        SqlLibraryOperatorTableFactory.INSTANCE.getOperatorTable(
                                SqlLibrary.STANDARD,
                                SqlLibrary.ORACLE
                        )
                )
        );

        MemSqlOperatorTable.registerFunctions(operatorTable);

        this.validator = SqlValidatorUtil.newValidator(
                operatorTable, catalogReader, typeFactory, SqlValidator.Config.DEFAULT
                        .withLenientOperatorLookup(config.lenientOperatorLookup())
                        .withSqlConformance(config.conformance())
                        .withDefaultNullCollation(config.defaultNullCollation())
                        .withIdentifierExpansion(true)
        );
        this.relConverter = createSqlToRelConverter();
        initExecutePlan();
    }

    private void registerTables(CalciteSchema rootSchema) {
        rootSchema.add("root", new ReflectiveSchema(new MemSchema()));

        SchemaPlus plus = rootSchema.plus();
        {
            Map<String, String> columns = Def.columns1;
            List<Pair<String, Class<?>>> nameTypes = columns.entrySet().stream().map(entry -> {
                String column = entry.getKey();
                Class<?> clz = JavaTypeResolve.resolveJavaType(entry.getValue());
                return Pair.<String, Class<?>>of(column, clz);
            }).collect(Collectors.toList());
            TableDefine tableDefine = new TableDefine("MOT_ORDER", nameTypes);
            plus.add(tableDefine.name(), new MemTable(tableDefine));
        }

        {
            Map<String, String> columns = Def.columns2;
            List<Pair<String, Class<?>>> nameTypes = columns.entrySet().stream().map(entry -> {
                String column = entry.getKey();
                Class<?> clz = JavaTypeResolve.resolveJavaType(entry.getValue());
                return Pair.<String, Class<?>>of(column, clz);
            }).collect(Collectors.toList());
            TableDefine tableDefine = new TableDefine("ORDER_ITEM", nameTypes);
            plus.add(tableDefine.name(), new MemTable(tableDefine));
        }
    }

    private void registerFunctions(CalciteSchema rootSchema) {
        SchemaPlus plus = rootSchema.plus();

    }

    private void initExecutePlan() throws SqlParseException {
        SqlNode sqlNode = SqlParser.create(sql, parserConfig).parseQuery();
        SqlNode validNode = validator.validate(sqlNode);
        RelNode logicPlan = relConverter.convertQuery(validNode, false, true).rel;
//        System.out.println(RelOptUtil.dumpPlan("[Logical plan]", logicPlan, SqlExplainFormat.TEXT, SqlExplainLevel.NON_COST_ATTRIBUTES))

        rowType = logicPlan.getRowType();

        final List<Pair<Integer, String>> fields =
                Pair.zip(ImmutableIntList.identity(rowType.getFieldCount()), rowType.getFieldNames());
        final RelCollation collation = logicPlan instanceof Sort
                ? ((Sort) logicPlan).collation
                : RelCollations.EMPTY;
        RelRoot root = new RelRoot(logicPlan, logicPlan.getRowType(), SqlKind.SELECT, fields,
                collation, new ArrayList<>());

        RelRoot root1 = this.optimize(root);

        this.bestPlan = root1.rel;

//        logicPlan = this.optimize(bestPlan);

        /*final RelShuttle shuttle = new RelHomogeneousShuttle() {
            @Override
            public RelNode visit(TableScan scan) {
                final RelOptTable table = scan.getTable();
                if (scan instanceof LogicalTableScan
                        && Bindables.BindableTableScan.canHandle(table)) {
                    // Always replace the LogicalTableScan with BindableTableScan
                    // because it's implementation does not require a "schema" as context.
                    return Bindables.BindableTableScan.create(scan.getCluster(), table);
                }
                return super.visit(scan);
            }
        };
        this.bestPlan = this.bestPlan.accept(shuttle);*/

//        System.out.println(RelOptUtil.dumpPlan("[Physical plan]", bestPlan, SqlExplainFormat.TEXT, SqlExplainLevel.NON_COST_ATTRIBUTES));
        // Obtain the executable plan
        Map<String, Object> internalParameters = new LinkedHashMap<>();
        executablePlan = EnumerableInterpretable.toBindable(
                internalParameters,
                null,
                (EnumerableRel) bestPlan,
                EnumerableRel.Prefer.ARRAY);
    }


    protected RelRoot optimize(RelRoot root) {
        final RelOptPlanner planner = root.rel.getCluster().getPlanner();
        final DataContext dataContext = createDataContext();
        planner.setExecutor(new RexExecutorImpl(dataContext));
        final RelTraitSet desiredTraits = getDesiredRootTraitSet(root);
        final RelNode rootRel4 = Programs.standard().run(
                planner, root.rel, desiredTraits, Collections.EMPTY_LIST, Collections.EMPTY_LIST);
        if (LOGGER.isTraceEnabled()) {
            LOGGER.trace("Plan after physical tweaks: {}",
                    RelOptUtil.toString(rootRel4, SqlExplainLevel.ALL_ATTRIBUTES));
        }
        return root.withRel(rootRel4);
    }

    protected RelTraitSet getDesiredRootTraitSet(RelRoot root) {
        // Make sure non-CallingConvention traits, if any, are preserved
        return root.rel.getTraitSet()
                .replace(EnumerableConvention.INSTANCE)
                .replace(root.collation)
                .simplify();
    }


    private SqlToRelConverter createSqlToRelConverter() {
        RelOptCluster cluster = newCluster();
        SqlToRelConverter.Config config = SqlToRelConverter.config().withTrimUnusedFields(true);
        ViewExpander expander = (rowType, queryString, schemaPath, viewPath) -> null;
        return new SqlToRelConverter(expander, validator, catalogReader, cluster, StandardConvertletTableExt.INSTANCE, config);
    }

    private RelOptCluster newCluster() {
        VolcanoPlanner planner = new VolcanoPlanner();
        addPlanRules(planner);
        planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
        planner.addRelTraitDef(RelCollationTraitDef.INSTANCE);
        planner.setTopDownOpt(false);

//        planner.addRelTraitDef(RelDistributionTraitDef.INSTANCE);
        return RelOptCluster.create(planner, new RexBuilder(typeFactory));
    }

    private void addPlanRules(final RelOptPlanner planner) {
        RelOptUtil.registerDefaultRules(planner, true, false);
    }

    private RelNode optimize(final RelNode logicPlan) {
        RelOptPlanner planner = relConverter.getCluster().getPlanner();
        planner.setRoot(planner.changeTraits(logicPlan, relConverter.getCluster().traitSet().replace(EnumerableConvention.INSTANCE)));
        return planner.findBestExp();
    }

    public Enumerable<Object[]> extract(final JSONObject data) {
        try {
            DataSetHolder.set(data);
            RelMetadataQueryBase.THREAD_PROVIDERS.set(JaninoRelMetadataProvider.DEFAULT);
            DataContext ctx = createDataContext();
            return executablePlan.bind(ctx);
        } finally {
            DataSetHolder.remove();
            RelMetadataQueryBase.THREAD_PROVIDERS.remove();
        }
    }

    public DataContext createDataContext() {
        return new DataContext() {
            @Override
            public SchemaPlus getRootSchema() {
                return rootSchema.plus();
            }

            @Override
            public JavaTypeFactory getTypeFactory() {
                return (JavaTypeFactory) typeFactory;
            }

            @Override
            public QueryProvider getQueryProvider() {
                return null;
            }

            @Override
            public Object get(String name) {
                return null;
            }
        };
    }

    public String getSql() {
        return sql;
    }

    public RelNode getRelNode() {
        return bestPlan;
    }

    public RelDataType getRowType() {
        return rowType;
    }

    public Bindable<Object[]> getExecutablePlan() {
        return executablePlan;
    }
}
